// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/api_aggregate.h>

#include <nebula/compute/exec.h>
#include <nebula/compute/function_internal.h>
#include <nebula/compute/registry.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula {

    namespace internal {
        template<>
        struct EnumTraits<compute::CountOptions::CountMode>
                : BasicEnumTraits<compute::CountOptions::CountMode, compute::CountOptions::ONLY_VALID,
                        compute::CountOptions::ONLY_NULL, compute::CountOptions::ALL> {
            static std::string name() { return "CountOptions::CountMode"; }

            static std::string value_name(compute::CountOptions::CountMode value) {
                switch (value) {
                    case compute::CountOptions::ONLY_VALID:
                        return "NON_NULL";
                    case compute::CountOptions::ONLY_NULL:
                        return "NULLS";
                    case compute::CountOptions::ALL:
                        return "ALL";
                }
                return "<INVALID>";
            }
        };

        template<>
        struct EnumTraits<compute::QuantileOptions::Interpolation>
                : BasicEnumTraits<compute::QuantileOptions::Interpolation,
                        compute::QuantileOptions::LINEAR, compute::QuantileOptions::LOWER,
                        compute::QuantileOptions::HIGHER, compute::QuantileOptions::NEAREST,
                        compute::QuantileOptions::MIDPOINT> {
            static std::string name() { return "QuantileOptions::Interpolation"; }

            static std::string value_name(compute::QuantileOptions::Interpolation value) {
                switch (value) {
                    case compute::QuantileOptions::LINEAR:
                        return "LINEAR";
                    case compute::QuantileOptions::LOWER:
                        return "LOWER";
                    case compute::QuantileOptions::HIGHER:
                        return "HIGHER";
                    case compute::QuantileOptions::NEAREST:
                        return "NEAREST";
                    case compute::QuantileOptions::MIDPOINT:
                        return "MIDPOINT";
                }
                return "<INVALID>";
            }
        };
    }  // namespace internal

    namespace compute {


        namespace internal {
            namespace {
                using ::nebula::internal::DataMember;
                static auto kScalarAggregateOptionsType = get_function_options_type<ScalarAggregateOptions>(
                        DataMember("skip_nulls", &ScalarAggregateOptions::skip_nulls),
                        DataMember("min_count", &ScalarAggregateOptions::min_count));
                static auto kCountOptionsType =
                        get_function_options_type<CountOptions>(DataMember("mode", &CountOptions::mode));
                static auto kModeOptionsType = get_function_options_type<ModeOptions>(
                        DataMember("n", &ModeOptions::n), DataMember("skip_nulls", &ModeOptions::skip_nulls),
                        DataMember("min_count", &ModeOptions::min_count));
                static auto kVarianceOptionsType = get_function_options_type<VarianceOptions>(
                        DataMember("ddof", &VarianceOptions::ddof),
                        DataMember("skip_nulls", &VarianceOptions::skip_nulls),
                        DataMember("min_count", &VarianceOptions::min_count));
                static auto kQuantileOptionsType = get_function_options_type<QuantileOptions>(
                        DataMember("q", &QuantileOptions::q),
                        DataMember("interpolation", &QuantileOptions::interpolation),
                        DataMember("skip_nulls", &QuantileOptions::skip_nulls),
                        DataMember("min_count", &QuantileOptions::min_count));
                static auto kTDigestOptionsType = get_function_options_type<TDigestOptions>(
                        DataMember("q", &TDigestOptions::q), DataMember("delta", &TDigestOptions::delta),
                        DataMember("buffer_size", &TDigestOptions::buffer_size),
                        DataMember("skip_nulls", &TDigestOptions::skip_nulls),
                        DataMember("min_count", &TDigestOptions::min_count));
                static auto kIndexOptionsType =
                        get_function_options_type<IndexOptions>(DataMember("value", &IndexOptions::value));
            }  // namespace
        }  // namespace internal

        ScalarAggregateOptions::ScalarAggregateOptions(bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kScalarAggregateOptionsType),
                  skip_nulls(skip_nulls),
                  min_count(min_count) {}

        constexpr char ScalarAggregateOptions::kTypeName[];

        CountOptions::CountOptions(CountMode mode)
                : FunctionOptions(internal::kCountOptionsType), mode(mode) {}

        constexpr char CountOptions::kTypeName[];

        ModeOptions::ModeOptions(int64_t n, bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kModeOptionsType),
                  n{n},
                  skip_nulls{skip_nulls},
                  min_count{min_count} {}

        constexpr char ModeOptions::kTypeName[];

        VarianceOptions::VarianceOptions(int ddof, bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kVarianceOptionsType),
                  ddof(ddof),
                  skip_nulls(skip_nulls),
                  min_count(min_count) {}

        constexpr char VarianceOptions::kTypeName[];

        QuantileOptions::QuantileOptions(double q, enum Interpolation interpolation,
                                         bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kQuantileOptionsType),
                  q{q},
                  interpolation{interpolation},
                  skip_nulls{skip_nulls},
                  min_count{min_count} {}

        QuantileOptions::QuantileOptions(std::vector<double> q, enum Interpolation interpolation,
                                         bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kQuantileOptionsType),
                  q{std::move(q)},
                  interpolation{interpolation},
                  skip_nulls{skip_nulls},
                  min_count{min_count} {}

        constexpr char QuantileOptions::kTypeName[];

        TDigestOptions::TDigestOptions(double q, uint32_t delta, uint32_t buffer_size,
                                       bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kTDigestOptionsType),
                  q{q},
                  delta{delta},
                  buffer_size{buffer_size},
                  skip_nulls{skip_nulls},
                  min_count{min_count} {}

        TDigestOptions::TDigestOptions(std::vector<double> q, uint32_t delta,
                                       uint32_t buffer_size, bool skip_nulls, uint32_t min_count)
                : FunctionOptions(internal::kTDigestOptionsType),
                  q{std::move(q)},
                  delta{delta},
                  buffer_size{buffer_size},
                  skip_nulls{skip_nulls},
                  min_count{min_count} {}

        constexpr char TDigestOptions::kTypeName[];

        IndexOptions::IndexOptions(std::shared_ptr<Scalar> value)
                : FunctionOptions(internal::kIndexOptionsType), value{std::move(value)} {}

        IndexOptions::IndexOptions() : IndexOptions(std::make_shared<NullScalar>()) {}

        constexpr char IndexOptions::kTypeName[];

        namespace internal {
            void register_aggregate_options(FunctionRegistry *registry) {
                        KCHECK_OK(registry->add_function_options_type(kScalarAggregateOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kCountOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kModeOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kVarianceOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kQuantileOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kTDigestOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kIndexOptionsType));
            }
        }  // namespace internal

        // ----------------------------------------------------------------------
        // Scalar aggregates

        turbo::Result<Datum> Count(const Datum &value, const CountOptions &options, ExecContext *ctx) {
            return call_function("count", {value}, &options, ctx);
        }

        turbo::Result<Datum> Mean(const Datum &value, const ScalarAggregateOptions &options,
                                  ExecContext *ctx) {
            return call_function("mean", {value}, &options, ctx);
        }

        turbo::Result<Datum> Product(const Datum &value, const ScalarAggregateOptions &options,
                                     ExecContext *ctx) {
            return call_function("product", {value}, &options, ctx);
        }

        turbo::Result<Datum> Sum(const Datum &value, const ScalarAggregateOptions &options,
                                 ExecContext *ctx) {
            return call_function("sum", {value}, &options, ctx);
        }

        turbo::Result<Datum> First(const Datum &value, const ScalarAggregateOptions &options,
                                   ExecContext *ctx) {
            return call_function("first", {value}, &options, ctx);
        }

        turbo::Result<Datum> Last(const Datum &value, const ScalarAggregateOptions &options,
                                  ExecContext *ctx) {
            return call_function("last", {value}, &options, ctx);
        }

        turbo::Result<Datum> MinMax(const Datum &value, const ScalarAggregateOptions &options,
                                    ExecContext *ctx) {
            return call_function("min_max", {value}, &options, ctx);
        }

        turbo::Result<Datum> Any(const Datum &value, const ScalarAggregateOptions &options,
                                 ExecContext *ctx) {
            return call_function("any", {value}, &options, ctx);
        }

        turbo::Result<Datum> All(const Datum &value, const ScalarAggregateOptions &options,
                                 ExecContext *ctx) {
            return call_function("all", {value}, &options, ctx);
        }

        turbo::Result<Datum> Mode(const Datum &value, const ModeOptions &options, ExecContext *ctx) {
            return call_function("mode", {value}, &options, ctx);
        }

        turbo::Result<Datum> Stddev(const Datum &value, const VarianceOptions &options,
                                    ExecContext *ctx) {
            return call_function("stddev", {value}, &options, ctx);
        }

        turbo::Result<Datum> Variance(const Datum &value, const VarianceOptions &options,
                                      ExecContext *ctx) {
            return call_function("variance", {value}, &options, ctx);
        }

        turbo::Result<Datum> Quantile(const Datum &value, const QuantileOptions &options,
                                      ExecContext *ctx) {
            return call_function("quantile", {value}, &options, ctx);
        }

        turbo::Result<Datum> TDigest(const Datum &value, const TDigestOptions &options,
                                     ExecContext *ctx) {
            return call_function("tdigest", {value}, &options, ctx);
        }

        turbo::Result<Datum> Index(const Datum &value, const IndexOptions &options, ExecContext *ctx) {
            return call_function("index", {value}, &options, ctx);
        }

    }  // namespace compute
}  // namespace nebula
